package com.llmops.core.node;

import com.llmops.core.Context;
import com.llmops.core.Edge;
import com.llmops.core.Node;
import com.llmops.core.NodeProperty;
import com.llmops.core.memory.MapDbMemoryStore;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.*;
import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;

public class LLMNode extends Node {
    @Data
    public static class LLMNodeProperty extends NodeProperty {

        private String apiUrl;
        private String apiKey;
        private String model;
        private Integer historyRecordCount;
        private String userPrompt;
        private String systemPrompt;
    }
    interface AssistantWithMemory{
        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
    }
    interface Assistant {
        String chat(@UserMessage String userMessage);
        @SystemMessage("{{systemPrompt}}")
        String chat(@UserMessage String userMessage,@V("systemPrompt")String systemPrompt);
    }

    private LLMNodeProperty property = new LLMNodeProperty();

    public LLMNode(String id, String title,Map<String, Object> config) {
        super(id, title,config);
        this.property.setApiUrl((String) config.getOrDefault("apiUrl", "https://api.deepseek.com/v1"));
        this.property.setApiKey((String) config.getOrDefault("apiKey","sk-0c3e6a9f9f7a4e63bb855290c544183c"));
        this.property.setModel((String) config.getOrDefault("model", "deepseek-chat"));
        this.property.setHistoryRecordCount((Integer) config.getOrDefault("historyRecordCount", 0));
        this.property.setUserPrompt((String) config.get("userPrompt"));
        this.property.setSystemPrompt((String) config.get("systemPrompt"));
    }

    @Override
    public void execute(Context context) {
        try {
            // 解析prompt中的变量
            String userPrompt;
            String systemPrompt;
            try {
                userPrompt = context.fillVariable(context,property.getUserPrompt());
                systemPrompt = context.fillVariable(context,property.getSystemPrompt());
            } catch (IllegalArgumentException e) {
                throw new RuntimeException("处理prompt变量时出错: " + e.getMessage(), e);
            }

            Integer sessionId = context.getVariable("sessionId");

            OpenAiChatModel chatModel = OpenAiChatModel.builder()
                    .baseUrl(property.getApiUrl())
                    .apiKey(property.getApiKey())
                    .modelName(property.getModel())
//                    .listeners(List.of(new ChatModelLogListener()))
                    .logRequests(true)
                    .logResponses(true)
                    .build();

            String answer;
            if(property.getHistoryRecordCount()>0){
                MapDbMemoryStore store = new MapDbMemoryStore();
                ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .maxMessages(property.getHistoryRecordCount())
                        .chatMemoryStore(store)
                        .build();
                AssistantWithMemory assistantWithMemory = AiServices.builder(AssistantWithMemory.class)
                        .chatModel(chatModel)
                        .chatMemoryProvider(chatMemoryProvider).build();
                answer = assistantWithMemory.chat(sessionId, userPrompt);
            }else {
                Assistant assistant = AiServices.builder(Assistant.class)
                        .chatModel(chatModel)
                        .build();
                answer = assistant.chat(userPrompt,systemPrompt);
            }
            context.setVariable(id + "_output", answer);
            logger.info(getId() + "_output={}", answer);
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("节点={},运行异常",id));
        }
    }



    @Override
    public String getType() {
        return "llmNode";
    }
}